import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from model import ImageClassifier
from utils import load_model

# 定义类别名称
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

def visualize_predictions(model, test_loader, device, num_images=9):
    model.eval()
    images, labels = next(iter(test_loader))
    images, labels = images.to(device), labels.to(device)
    
    # 随机选择一些图片
    indices = torch.randperm(images.size(0))[:num_images]
    images = images[indices]
    labels = labels[indices]

    # 获取预测结果
    with torch.no_grad():
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)

    # 显示图片和预测结果
    plt.figure(figsize=(10, 10))
    for i in range(num_images):
        plt.subplot(3, 3, i + 1)
        img = images[i] / 2 + 0.5  # 反归一化
        img = img.cpu().numpy().transpose((1, 2, 0))
        plt.imshow(img)
        plt.title(f"Pred: {classes[predicted[i]]}\nTrue: {classes[labels[i]]}")
        plt.axis('off')
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # 加载模型
    model = ImageClassifier().to(device)
    load_model(model, "image_classifier.pth")

    # 加载测试数据
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    test_dataset = CIFAR10(root='./data', train=False, download=True, transform=transform)
    test_loader = DataLoader(test_dataset, batch_size=100, shuffle=True)

    # 可视化预测结果
    visualize_predictions(model, test_loader, device)